from contextlib import contextmanager
from datetime import datetime, date

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy
from flask_sqlalchemy import BaseQuery

from sqlalchemy import SmallInteger, Column, Integer

from apps.libs.error_code import NotFound


class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e


class Query(BaseQuery):
    def filter_by(self, **kwargs):
        if 'status' not in kwargs.keys():
            kwargs['status'] = 1
        return super(Query, self).filter_by(**kwargs)

    def get_or_404(self, ident):
        rv = self.get(ident)
        if not rv:
            raise NotFound()
        return rv

    def first_or_404(self):
        rv = self.first()
        if not rv:
            raise NotFound()
        return rv


db = SQLAlchemy(query_class=Query)


class Base(db.Model):
    __abstract__ = True

    status = Column(SmallInteger, default=1)
    create_time = Column('create_time', Integer)

    def __init__(self):
        self.create_time = int(datetime.now().timestamp())

    # 实现此方法允许对象以字典的方式调用属性self['']
    def __getitem__(self, item):
        return getattr(self, item)

    def set_attrs(self, attrs_dict):
        for key, value in attrs_dict.item():
            if hasattr(self, key):
                setattr(self, key, value)

    @property
    def create_datetime(self):
        if self.create_time:
            return date.fromtimestamp(self.create_time)
        else:
            return 'Null'

    def delete(self):
        self.status = 0

    # 解决有些接口不需要某几个字段的特殊需求，可以使用对象实例实现hide方法，传入指定隐藏的字段
    def hide(self, *keys):
        for key in keys:
            self.fields.remove(key)
        return self

    def append(self, *keys):
        for key in keys:
            self.field.append(key)
        return self
